import torch
from torch.utils.data import Dataset,DataLoader
from torchsummary import summary
import torch.nn as nn
from torch.optim import SGD
import matplotlib.pyplot as plt

x = [[1,2],[3,4],[5,6],[7,8]]
y = [[3],[7],[11],[15]]

X = torch.tensor(x).float()
Y = torch.tensor(y).float()

# device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = 'cpu'
X = X.to(device)
Y = Y.to(device)

class MyDataset(Dataset):
    def __init__(self,x,y):
        self.x = torch.tensor(x).float().to(device)
        self.y = torch.tensor(y).float().to(device)

    def __getitem__(self, item):
        return self.x[item],self.y[item]

    def __len__(self):
        return len(self.x)


ds = MyDataset(X,Y)
dl = DataLoader(ds,batch_size=2,shuffle=True)

model = nn.Sequential(
    nn.Linear(2,8),
    nn.ReLU(),
    nn.Linear(8,1)
).to(device)

# print(summary(model))

loss_fun = nn.MSELoss()
opt = SGD(model.parameters(),lr = 0.001)
loss_history = []
for _ in range(50):
    for ix,iy in dl:
        opt.zero_grad()
        loss_value = loss_fun(model(ix),iy)
        loss_value.backward()
        opt.step()
        loss_history.append(loss_value.detach().numpy())

plt.figure(figsize=(10, 10))
plt.plot(loss_history)
plt.show()